package com.qyx.pmpucat.util;

import org.apache.commons.math3.linear.RealMatrix;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class SelectNextItemMfi {

    private static final double D = 1.0;

    // 计算项目反应理论 (IRT) 中的 Pi 矩阵及其导数
    public static PiResult Pi(double th, RealMatrix it, double D) {
        int nItems = it.getRowDimension();
        int maxCategories = it.getColumnDimension(); // 减去区分度参数

        double[][] prov = new double[nItems][maxCategories];
        double[][] prov1 = new double[nItems][maxCategories];
        double[][] prov2 = new double[nItems][maxCategories];
        double[][] prov3 = new double[nItems][maxCategories];

        for (int i = 0; i < nItems; i++) {
            double aj = it.getEntry(i, 0);
            double[] bj = Arrays.stream(it.getRow(i)).skip(1).filter(b -> !Double.isNaN(b)).toArray();

            double[] Pjs = new double[bj.length + 2];
            Pjs[0] = 1.0; // 边界条件：0分及0分以上的概率为1
            for (int j = 0; j < bj.length; j++) {
                double ej = Math.exp(D * aj * (th - bj[j]));
                Pjs[j + 1] = ej / (1 + ej); // 根据公式计算中间概率
            }
            Pjs[Pjs.length - 1] = 0.0; // 边界条件：大于最大分值的概率为0

            double[] dPjs = new double[Pjs.length];
            double[] d2Pjs = new double[Pjs.length];
            double[] d3Pjs = new double[Pjs.length];

            for (int j = 0; j < Pjs.length; j++) {
                dPjs[j] = D * aj * Pjs[j] * (1 - Pjs[j]);
                d2Pjs[j] = D * aj * (dPjs[j] - 2 * Pjs[j] * dPjs[j]);
                d3Pjs[j] = D * aj * (d2Pjs[j] - 2 * Math.pow(dPjs[j], 2) - 2 * Pjs[j] * d2Pjs[j]);
            }

            // 确保不会超出Pjs数组的边界
            int maxValidIndex = Math.min(maxCategories, Pjs.length - 1);
            for (int j = 0; j < maxValidIndex; j++) {
                prov[i][j] = Pjs[j] - Pjs[j + 1];
                prov1[i][j] = dPjs[j] - dPjs[j + 1];
                prov2[i][j] = d2Pjs[j] - d2Pjs[j + 1];
                prov3[i][j] = d3Pjs[j] - d3Pjs[j + 1];
            }
        }

        return new PiResult(prov, prov1, prov2, prov3);
    }

    // 计算 Fisher 信息量 Ii
    public static Map<String, double[]> Ii(double th, RealMatrix it, double D) {
        PiResult pr = Pi(th, it, D);
        double[] Ii = new double[it.getRowDimension()];
        double[] dIi = new double[it.getRowDimension()];
        double[] d2Ii = new double[it.getRowDimension()];

        for (int i = 0; i < it.getRowDimension(); i++) {
            // 获取题目原始参数
            double[] rawParams = it.getRow(i);
            // 提取有效难度参数（跳过区分度参数aj）
            double[] bj = Arrays.stream(rawParams)
                    .skip(1)
                    .filter(b -> !Double.isNaN(b))
                    .toArray();

            double pr0 = 0, pr1 = 0, pr2 = 0;
            for (int j = 0; j < pr.Pi[i].length; j++) {
                // 新增判断：当难度参数为0时跳过计算
                if (j < bj.length && bj[j] == 0.0) continue;

                if (!Double.isNaN(pr.dPi[i][j]) && pr.Pi[i][j] > 0) {
                    pr0 += Math.pow(pr.dPi[i][j], 2) / pr.Pi[i][j];// 计算信息量
                    pr1 += 2 * pr.dPi[i][j] * pr.d2Pi[i][j] / pr.Pi[i][j] - Math.pow(pr.dPi[i][j], 3) / Math.pow(pr.Pi[i][j], 2);
                    pr2 += (2 * Math.pow(pr.d2Pi[i][j], 2) + 2 * pr.dPi[i][j] * pr.d3Pi[i][j]) / pr.Pi[i][j] - 5 * Math.pow(pr.dPi[i][j], 2) * pr.d2Pi[i][j] / Math.pow(pr.Pi[i][j], 2) + 2 * Math.pow(pr.dPi[i][j], 4) / Math.pow(pr.Pi[i][j], 3);
                }
            }
            Ii[i] = pr0; // 信息量；信息量越大，越有信息
            dIi[i] = pr1;
            d2Ii[i] = pr2;
        }

        Map<String, double[]> res = new HashMap<>();
        res.put("Ii", Ii);
        res.put("dIi", dIi);
        res.put("d2Ii", d2Ii);
        return res;
    }

    // 最大信息量选题函数
    public static Map<String, Object> selectItemMFI(double theta, RealMatrix itemBank, Set<Integer> OUT, int randomesque, double D) {
        // 计算每个题目在当前能力估计值 theta 下的信息量
        Map<String, double[]> infoMap = Ii(theta, itemBank, D);
        double[] info = infoMap.get("Ii");
        //索引/需要-1
        Set<Integer> collect = OUT.stream().map(i -> i = i - 1).collect(Collectors.toSet());
        // 创建题目索引与信息量的映射，并移除已作答的题目
        List<Pair<Integer, Double>> itemInfoPairs = IntStream.range(0, info.length)
                .filter(i -> !collect.contains(i)) // 排除已作答的题目
                .mapToObj(i -> new Pair<>(i, info[i]))
                .collect(Collectors.toList());

        // 按照信息量排序
        // itemInfoPairs.sort((p1, p2) -> Double.compare(p2.getValue(), p1.getValue()));
        itemInfoPairs.sort(Comparator.comparingDouble((Pair<Integer, Double> p) -> p.getValue()).reversed());

        // 确定随机选题范围
        int nrIt = Math.min(randomesque, itemInfoPairs.size());
        List<Integer> topItems = itemInfoPairs.stream()
                .limit(nrIt)
                .map(Pair::getKey)
                .collect(Collectors.toList());

        // 随机从范围中选题
        Random rand = new Random();
        Integer select = topItems.isEmpty() ? null : topItems.get(rand.nextInt(topItems.size()));

        // 返回选题结果
        Map<String, Object> res = new HashMap<>();
        if (select != null) {
            res.put("item", select + 1);
            res.put("par", Arrays.toString(itemBank.getRow(select)));
            res.put("info", info[select]);
        } else {
            res.put("item", "No items available");
            res.put("par", "N/A");
            res.put("info", "N/A");
        }
        res.put("criterion", "MFI");
        res.put("randomesque", randomesque);
        res.put("name", null);

        return res;
    }

    // 定义一个内部类 PiResult，用于存储 Pi 及其导数的结果
    public static class PiResult {
        public double[][] Pi;  // 概率矩阵
        public double[][] dPi;  // 一阶导数矩阵
        public double[][] d2Pi;  // 二阶导数矩阵
        public double[][] d3Pi;  // 三阶导数矩阵

        public PiResult(double[][] Pi, double[][] dPi, double[][] d2Pi, double[][] d3Pi) {
            this.Pi = Pi;
            this.dPi = dPi;
            this.d2Pi = d2Pi;
            this.d3Pi = d3Pi;
        }
    }

    // 辅助类，用于存储键值对
    private static class Pair<K, V> {
        private final K key;
        private final V value;

        public Pair(K key, V value) {
            this.key = key;
            this.value = value;
        }

        public K getKey() {
            return key;
        }

        public V getValue() {
            return value;
        }
    }
}